package etxt2db.api;

import java.io.File;
import java.io.IOException;
import java.text.ParseException;
import java.util.HashSet;
import java.util.List;
import java.util.Set;


import edu.cmu.minorthird.text.MutableTextLabels;
import edu.cmu.minorthird.text.learn.SpanFeatureExtractor;
import edu.cmu.minorthird.ui.Recommended;
import etxt2db.annotators.CRFClassificationModel;
import etxt2db.annotators.DictionaryClassificationModel;
import etxt2db.annotators.HMMClassificationModel;
import etxt2db.annotators.MEMMClassificationModel;
import etxt2db.annotators.MinorthirdCRFClassificationModel;
import etxt2db.annotators.MinorthirdMEMMClassificationModel;
import etxt2db.annotators.MyDictionaryClassificationModel;
import etxt2db.annotators.NaiveDictionaryClassificationModel;
import etxt2db.annotators.RegexClassificationModel;
import etxt2db.annotators.SVMClassificationModel;
import etxt2db.corpus.MyCompleteTextBaseLoader;
import etxt2db.corpus.MyTextBaseLoader;

/**
 * <p>
 * This class represents an object that contains functions to train an Information Extraction
 * classifier given a training corpus.
 * </p>
 * <p>
 * The training corpus is formated as a plain text with XML tags representing the classification of
 * text segments. Following we present an example:
 * </p>
 * 
 * <blockquote><pre>
 * CENTER FOR INNOVATION IN LEARNING (CIL)<br/>
 * EDUCATION SEMINAR SERIES<br/><br/>
 * 
 * "Using a Cognitive Architecture to Design Instructions"<br/><br/>
 * 
 * &lt;speaker&gt;Joe Mertz&lt;/speaker&gt;<br/>
 * Center for Innovation in Learning, CMU<br/><br/>
 * 
 * Friday, February 17<br/>
 * &lt;stime&gt;12:00pm&lt;/stime&gt;-&lt;etime&gt;1:00pm&lt;/etime&gt;<br/>
 * &lt;location&gt;Student Center Room 207&lt;/location&gt; (CMU)<br/>
 *  ABSTRACT: In my talk, I will describe how a cognitive model was used
 *  as a simulated student to help design instructions for training
 *  circuit board assemblers.  The model was built in the Soar cognitive
 *  architecture, and was initially endowed with only minimal prerequisite
 *  knowledge of the task, and an ability to learn instructions.  Lessons
 *  for teaching expert assembly skills were developed by iteratively
 *  drafting and testing instructions on the simulated student.<br/><br/>
 * 
 *  Please direct questions to Pamela Yocca at 268-7675.
 *  </pre></blockquote>
 *
 * @author Gon�alo Sim�es
 */


public class ClassificationModelCreator {
	
	public enum MLTechnique {HMM, SVM, CRF, MEMM}
	
	private Set<Set<String>> compatibleAttributes = new HashSet<Set<String>>();
	private boolean docPerFile = true;
	private String beginTagString = "<";
	private String endTagString = ">";
	
	public ClassificationModelCreator(){
		this.compatibleAttributes = new HashSet<Set<String>>();
	}
	
	/**
     * Function to train a classification model based on machine learning
     * @param file A File containing the Training corpus
     * @param technique The technique with which the model will be trained
     * @param attributes The list containing the attributes for this classification
     * @return ClassificationModel object able to classify text segments
     */
	
	public ClassificationModel trainMachineLearningModel(File file, MLTechnique technique, List<String> attributes) 
		throws IOException, ParseException{
		return trainMachineLearningModel(file, technique, attributes, new Recommended.TokenFE());
	}
	
	/**
     * Function to train a classification model based on machine learning
     * @param file A File containing the Training corpus
     * @param technique The technique with which the model will be trained
     * @param attributes The list containing the attributes for this classification
     * @param tokenFE The Feature Extractor you want to use
     * @return ClassificationModel object able to classify text segments
     */
	
	public ClassificationModel trainMachineLearningModel(File file, MLTechnique technique, List<String> attributes, SpanFeatureExtractor tokenFE) 
		throws IOException, ParseException{
		MutableTextLabels trainingFile;
		switch (technique) {
			case HMM:
				trainingFile = loadTrainingDocument(file, attributes);
				return new HMMClassificationModel(trainingFile,compatibleAttributes);
			case SVM:
				trainingFile = loadTrainingDocument(file, attributes);
				return new SVMClassificationModel(trainingFile, tokenFE,compatibleAttributes);
			case CRF:
				trainingFile = loadTrainingDocument(file, attributes);
				return new MinorthirdCRFClassificationModel(trainingFile, tokenFE,compatibleAttributes);
				//return new MalletCRFAnnotator(trainingFile,attributes,tokenFE);
			case MEMM: 
				trainingFile = loadTrainingDocument(file, attributes);
				return new MinorthirdMEMMClassificationModel(trainingFile, tokenFE,compatibleAttributes);
				//return new MalletMEMMAnnotator(trainingFile,attributes,tokenFE);
			default:
				trainingFile = loadTrainingDocument(file, attributes);
				return new HMMClassificationModel(trainingFile,compatibleAttributes);
		}
	}
	
	/**
     * Function to train a classification model based on regular expressions
     * @param regex A String with the regular expression that matches the class to classify
     * @param type A String with the name of the class to classify
     * @return ClassificationModel object able to classify text segments
     */
	public ClassificationModel createRegexClassificationModel(String regex, String type) {
		return RegexClassificationModel.getOptimizedClassificationModel(regex,type);
	}
	
	/**
     * Function to train a classification model based on a dictionary
     * @param dictionaryFile A File with pointing to the dictionary
     * @return ClassificationModel object able to classify text segments
     */
	public ClassificationModel createDictionaryClassificationModel(File dictionaryFile) throws IOException, ParseException {
		return new DictionaryClassificationModel(dictionaryFile);
	}
	
	/**
     * Function to train a classification model based on a dictionary
     * @param dictionaryFile A File with pointing to the dictionary
     * @param isCaseSensitive boolean value indicating whether the dictionary is case sensitive
     * @return ClassificationModel object able to classify text segments
     */
	public ClassificationModel createDictionaryClassificationModel(File dictionaryFile, boolean isCaseSensitive) throws IOException, ParseException {
		return new DictionaryClassificationModel(dictionaryFile,isCaseSensitive);
	}
	
	/**
     * Function to train a classification model based on a dictionary
     * @param dictionaryFile A File with pointing to the dictionary
     * @return ClassificationModel object able to classify text segments
     */
	public ClassificationModel createNaiveDictionaryClassificationModel(File dictionaryFile) throws IOException, ParseException {
		return new NaiveDictionaryClassificationModel(dictionaryFile);
	}
	
	/**
     * Function to train a classification model based on a dictionary
     * @param dictionaryFile A File with pointing to the dictionary
     * @param isCaseSensitive boolean value indicating whether the dictionary is case sensitive
     * @return ClassificationModel object able to classify text segments
     */
	public ClassificationModel createNaiveDictionaryClassificationModel(File dictionaryFile, boolean isCaseSensitive) throws IOException, ParseException {
		return new NaiveDictionaryClassificationModel(dictionaryFile,isCaseSensitive);
	}
	
	private MutableTextLabels loadTrainingDocument(File file, List<String> attributes) throws IOException, ParseException{
		MyCompleteTextBaseLoader baseLoader = new MyCompleteTextBaseLoader(attributes);
		if(!docPerFile){
			baseLoader.setDocumentStyle(MyTextBaseLoader.DOC_PER_LINE);
		}
		baseLoader.setBeginTag(beginTagString);
		baseLoader.setEndTag(endTagString);
		baseLoader.load(file);
		compatibleAttributes = baseLoader.getCompatiblePartition();
		MutableTextLabels labels = baseLoader.getLabels();
		return labels;
	}
	
	public void setBeginTag(String str){
		this.beginTagString = str;
	}
	
	public void setEndTag(String str){
		this.endTagString = str;
	}
	
	public void setDocPerFile(boolean docPerFile){
		this.docPerFile = docPerFile;
	}
}
